"""
Model Factory for creating different models for data assimilation
"""
from Tools.AssimilationSystem import NumericalModel
from Tools.barotropic_model import BarotropicModel

class ModelFactory:
    """
    Factory class for creating different numerical models
    """
    
    @staticmethod
    def create_model(model_type: str, state_size=40):
        """
        Create a model based on the specified type
        
        Parameters:
        -----------
        model_type : str
            Type of model to create ("BAROTROPIC" or "LINEAR")
        state_size : int
            Size of the state vector (used for LINEAR model)
            
        Returns:
        --------
        model : NumericalModel or BarotropicModel
            The created model
            
        Raises:
        -------
        ValueError
            If an invalid model_type is specified
        """
        if model_type not in ("BAROTROPIC", "LINEAR"):
            raise ValueError(f"Invalid model type: {model_type}. Must be 'BAROTROPIC' or 'LINEAR'")
        if model_type == "BAROTROPIC":
            # Calculate grid dimensions from state_size
            nx = next(i for i in range(int(state_size**0.5), 0, -1) 
                    if state_size % i == 0)
            ny = state_size // nx
            print(f"nx: {nx}, ny: {ny}")
            return BarotropicModel(nx=nx, ny=ny)
        else:
            return NumericalModel(state_size)
